import requests
import json
import tempfile
import os
import PyPDF2
from typing import Dict, Optional
import config


class LLMProcessor:
    def __init__(self):
        self.api_key = config.DEEPSEEK_API_KEY
        self.api_url = "https://api.deepseek.com/v1/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def _extract_text_from_pdf(self, pdf_path: str) -> str:
        """Extract text from a PDF file."""
        text = ""
        with open(pdf_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            for page in reader.pages:
                text += page.extract_text() + "\n"
        return text

    def process_paper(self, title: str, abstract: str, authors: str,
                      paper_file: Optional[str] = None, request_id: str = None) -> Dict:
        """Process paper information using DeepSeek API and return structured poster content."""

        # Extract text from PDF if provided
        paper_text = ""
        if paper_file and os.path.exists(paper_file):
            paper_text = self._extract_text_from_pdf(paper_file)

        # Construct prompt for the LLM
        prompt = f"""
        Create structured content for an academic poster based on the following paper:

        Title: {title}
        Authors: {authors}
        Abstract: {abstract}

        Additional Paper Content:
        {paper_text[:2000] if paper_text else "Not provided"}

        Please extract and organize the following elements for a poster:
        1. Main title
        2. Authors with affiliations
        3. Introduction/Background (2-3 key points)
        4. Research question or objectives
        5. Methodology (brief)
        6. Key results (3-5 points with potential visualizations)
        7. Conclusions and implications
        8. References (up to 5 key ones)
        9. Contact information

        Also suggest 2-3 key visualizations that would effectively represent the research.
        Format the response as a structured JSON.
        """

        # Call DeepSeek API
        response = requests.post(
            self.api_url,
            headers=self.headers,
            json={
                "model": config.LLM_MODEL,
                "messages": [
                    {"role": "system",
                     "content": "You are an expert academic researcher and designer who specializes in creating academic posters."},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.3,
                "max_tokens": 2000
            }
        )

        if response.status_code == 200:
            try:
                content = response.json()["choices"][0]["message"]["content"]
                # Parse the JSON content returned by the LLM
                structured_content = json.loads(content)
                return structured_content
            except json.JSONDecodeError:
                # If the LLM didn't return proper JSON, attempt to extract structured data
                return self._extract_structured_data(content)
        else:
            raise Exception(f"API error: {response.status_code} - {response.text}")

    def _extract_structured_data(self, text: str) -> Dict:
        """Extract structured data from text if JSON parsing fails."""
        # Implement fallback parsing logic
        sections = {}
        current_section = None

        for line in text.split("\n"):
            line = line.strip()
            if not line:
                continue

            if line.startswith("#") or line.endswith(":"):
                # This looks like a section header
                current_section = line.replace("#", "").replace(":", "").strip().lower()
                sections[current_section] = []
            elif current_section:
                sections[current_section].append(line)

        # Convert to a more structured format
        result = {
            "title": sections.get("main title", [""])[0] if "main title" in sections else "",
            "authors": sections.get("authors", []),
            "introduction": sections.get("introduction", []) or sections.get("background", []),
            "objectives": sections.get("research question", []) or sections.get("objectives", []),
            "methodology": sections.get("methodology", []),
            "results": sections.get("key results", []) or sections.get("results", []),
            "conclusions": sections.get("conclusions", []) or sections.get("implications", []),
            "references": sections.get("references", []),
            "contact": sections.get("contact information", []),
            "visualizations": sections.get("visualizations", []) or sections.get("figures", [])
        }

        return result